package pers.lls.hds.tree;

//https://blog.csdn.net/weixin_40037053/article/details/89947885
//https://www.cnblogs.com/lycroseup/p/7324229.html
public class RBTree<T extends Comparable<T>, D> {

    private RBNode<T, D> root;
    private static final Boolean RED = false;
    private static final Boolean BLACK = true;

    public class RBNode<T extends Comparable<T>, D> {

        private Boolean color;
        private T key;
        private D data;
        private RBNode<T, D> parent;
        private RBNode leftChild;
        private RBNode rightChild;


        public RBNode(Boolean col, T key, D data, RBNode paret, RBNode leftChild, RBNode rightChild) {
            this.color = col;
            this.key = key;
            this.data = data;
            this.parent = parent;
            this.leftChild = leftChild;
            this.rightChild = rightChild;

        }

    }

    public RBNode<T, D> parentOf(RBNode<T, D> node) {
        if (node != null) {
            return node.parent;

        }

        return null;

    }

    public Boolean colorOf(RBNode<T, D> node) {
        if (node != null) {
            return node.color;

        }
        return BLACK;

    }

    public void setParent(RBNode<T, D> node, RBNode<T, D> parent) {
        if (node != null) {
            node.parent = parent;
        }

    }

    public void setColor(RBNode<T, D> node, Boolean color) {
        if (node != null) {
            node.color = color;

        }

    }

    public Boolean isRed(RBNode<T, D> node) {
        return (node != null && node.color == RED) ? true : false;
    }

    public Boolean isBlack(RBNode<T, D> node) {
        return !isRed(node);

    }

    public void setRed(RBNode<T, D> node) {
        if (node != null) {
            node.color = RED;
        }
    }

    public void setBlack(RBNode<T, D> node) {
        if (node != null) {
            node.color = BLACK;
        }

    }

    public D get(T key) {
        RBNode node = search(key, root);
        return node == null ? null : (D) node.data;
    }

    public RBNode<T, D> search(T key, RBNode<T, D> node) {

        if (node != null) {
            //查找的过程，就是一直递归比较到叶子为止
            int com = key.compareTo(node.key);
            if (com < 0) {
                return search(key, node.leftChild);
            } else if (com > 0) {
                return search(key, node.rightChild);
            } else {
                return node;
            }

        }
        return null;


    }

    public RBNode<T, D> min(RBNode<T, D> node) {

        //一直往左走，最左端的就是最小值，这是二叉树的性质
        if (node.leftChild == null) {
            return node;
        }
        while (node.leftChild != null) {
            node = node.leftChild;
        }

        return node;
    }

    //寻找待删节点的后继节点
    public RBNode successor(RBNode<T, D> node) {

        if (node.rightChild != null) {
            return min(node.rightChild);
        }
        //下面这里是不会进入的，因为只有node的两个孩子都不为null时才会进入这个方法
        RBNode<T, D> y = node.parent;
        while ((y != null) && (y.rightChild == node)) {
            node = y;
            y = y.parent;

        }
        return y;

    }

    //左旋
    public void leftRonate(RBNode<T, D> x) {
        RBNode<T, D> y = x.rightChild;
        if (y.leftChild != null) {
            y.leftChild.parent = x;
        }
        x.rightChild = y.leftChild;
        y.leftChild = x;
        y.parent = x.parent;
        if (x.parent != null) {
            if (x.parent.leftChild == x) {
                x.parent.leftChild = y;
            } else {
                x.parent.rightChild = y;
            }
        } else {
            this.root = y;
        }
        x.parent = y;
    }

    //右旋
    public void rightRonate(RBNode<T, D> x) {
        RBNode<T, D> y = x.leftChild;
        if (y.rightChild != null) {
            y.rightChild.parent = x;
        }
        y.parent = x.parent;
        x.leftChild = y.rightChild;
        y.rightChild = x;
        if (x.parent != null) {
            if (x.parent.leftChild == x) {
                x.parent.leftChild = y;
            } else {
                x.parent.rightChild = y;
            }
        } else {
            this.root = y;
        }
        x.parent = y;
    }

    public void balanceInsertion(RBNode<T, D> node) {
        RBNode<T, D> parent, gparent;
        while (((parent = parentOf(node)) != null) && isRed(parent)) {
            gparent = parentOf(parent);
            if (gparent.leftChild == parent) {
                RBNode<T, D> uncle = gparent.rightChild;
                if (isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    node = gparent;
                    continue;
                } else {
                    if (parent.rightChild == node) {
                        //假如当前节点是父亲的右孩子
                        //父亲要左旋
                        leftRonate(parent);
                        RBNode<T, D> temp = node;
                        node = parent;
                        parent = temp;

                    }
                    setBlack(parent);
                    setRed(gparent);
                    rightRonate(gparent);
                }

            } else {
                RBNode<T, D> uncle = gparent.leftChild;
                if (isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    node = gparent;
                    continue;
                } else {
                    if (parent.leftChild == node) {
                        rightRonate(parent);
                        RBNode<T, D> temp = node;
                        node = parent;
                        parent = temp;
                    }
                    setBlack(parent);
                    setRed(gparent);
                    leftRonate(gparent);
                }
            }
        }
        if (root == node) {
            setBlack(node);
        }
    }

    public void balanceDeletion(RBNode<T, D> node, RBNode<T, D> parent) {
        RBNode<T, D> other;
        while (isBlack(node) && node != this.root) {
            if (parent.leftChild == node) {
                other = parent.rightChild;
                if (isRed(other)) {
                    setRed(parent);
                    setBlack(other);
                    leftRonate(parent);
                    continue;
                } else {
                    if (isBlack(other.leftChild) && isBlack(other.rightChild)) {
                        setRed(other);
                        node = parent;
                        parent = parentOf(node);
                    } else if (isRed(other.leftChild) && isBlack(other.rightChild)) {
                        setRed(other);
                        setBlack(other.leftChild);
                        rightRonate(other);
                    } else if (isRed(other.rightChild)) {
                        setColor(other, colorOf(parent));
                        setBlack(parent);
                        setBlack(other.rightChild);
                        leftRonate(parent);
                        break;
                    }
                }
            } else {
                other = parent.leftChild;
                if (isRed(other)) {
                    setBlack(other);
                    setRed(parent);
                    rightRonate(parent);
                    continue;
                } else {
                    if (isBlack(other.leftChild) && isBlack(other.rightChild)) {
                        setRed(other);
                        node = parent;
                        parent = parentOf(node);
                    } else if (isRed(other.rightChild) && isBlack(other.leftChild)) {
                        setRed(parent);
                        setBlack(other.rightChild);
                        leftRonate(other);
                    } else if (isRed(other.leftChild)) {
                        setColor(other, colorOf(parent));
                        setBlack(parent);
                        setBlack(other.leftChild);
                        rightRonate(parent);
                        break;
                    }
                }
            }
        }
        setBlack(node);
    }

    private void insertNode(T key, D data) {

        int com;
        RBNode<T, D> x = this.root;
        RBNode<T, D> y = null;

        while (x != null) {
            y = x;
            com = key.compareTo(x.key);

            if (com == 0) {
                x.data = data;
                return;
            }

            if (com < 0) {
                x = x.leftChild;
            } else {
                x = x.rightChild;
            }
        }
        RBNode<T, D> node = new RBNode<>(RED, key, data, null, null, null);
        node.parent = y;

        if (y != null) {
            com = node.key.compareTo(y.key);
            if (com < 0) {
                y.leftChild = node;
            } else {
                y.rightChild = node;
            }
        } else {
            this.root = node;
        }
        balanceInsertion(node);
    }

    public void insert(T key, D data) {
        insertNode(key, data);
    }

    private void delete(RBNode<T, D> node) {

        RBNode<T, D> child, parent, replace;
        Boolean color;
        if (node.leftChild != null && node.rightChild != null) {
            replace = successor(node);
            parent = parentOf(replace);
            child = replace.rightChild;
            color = colorOf(replace);
            if (node == parentOf(replace)) {
                parent = replace;
            } else {
                if (child != null) {
                    setParent(child, parentOf(replace));
                }
                replace.parent.leftChild = child;
                replace.rightChild = node.rightChild;
                setParent(node.rightChild, replace);
            }
            setParent(replace, parentOf(node));
            replace.leftChild = node.leftChild;
            setParent(node.leftChild, replace);
            setColor(replace, colorOf(node));
            if (parentOf(node) != null) {
                if (node.parent.leftChild == node) {
                    node.parent.leftChild = replace;
                } else {
                    node.parent.rightChild = replace;
                }
            } else {
                this.root = replace;
            }
            if (color == BLACK) {
                balanceDeletion(child, parent);
            }

        } else {
            if (node.leftChild != null) {
                replace = node.leftChild;
            } else {
                replace = node.rightChild;
            }
            parent = parentOf(node);

            if (parent != null) {
                if (parent.leftChild == node) {
                    parent.leftChild = replace;
                } else {
                    parent.rightChild = replace;
                }
            } else {
                this.root = replace;
            }
            setParent(replace, parent);
            color = colorOf(node);
            child = replace;
            if (color == BLACK) {
                balanceDeletion(child, parent);
            }
        }
    }

    public void delete(T key) {
        RBNode<T, D> node;
        if ((node = search(key, this.root)) != null) {
            delete(node);
        }
    }

    //前序遍历
    private void preOrder(RBNode<T, D> node) {
        if (node != null) {

            System.out.print(node.key + " ");
            preOrder(node.leftChild);
            preOrder(node.rightChild);

        }
    }

    public void preOrder() {
        preOrder(this.root);
    }

    //中序遍历
    private void inOrder(RBNode<T, D> node) {
        if (node != null) {
            inOrder(node.leftChild);
            System.out.print(node.key + " ");
            inOrder(node.rightChild);

        }

    }

    public void inOrder() {
        inOrder(this.root);
    }

    //后序遍历
    private void postOrder(RBNode<T, D> node) {
        if (node != null) {
            postOrder(node.leftChild);
            postOrder(node.rightChild);
            System.out.print(node.key + " ");

        }

    }

    public void postOrder() {
        postOrder(this.root);
    }


    public static void main(String[] args) {
        RBTree tree = new RBTree();
        tree.insert(22,22);
        tree.insert(13,13);
        tree.insert(15,15);
        tree.insert(17,17);
        tree.insert(25,25);
        tree.insert(27,27);
        tree.insert(1,1);
        tree.insert(6,6);
        tree.insert(8,8);
        tree.insert(11,11);
        tree.show(tree.root);
    }


    // 用于获得树的层数
    public int getTreeDepth(RBNode root) {
        return root == null ? 0 : (1 + Math.max(getTreeDepth(root.leftChild), getTreeDepth(root.rightChild)));
    }

    private void writeArray(RBNode currNode, int rowIndex, int columnIndex, String[][] res, int treeDepth) {
        // 保证输入的树不为空
        if (currNode == null) return;
        // 先将当前节点保存到二维数组中
        res[rowIndex][columnIndex] = (currNode.color ? "b" : "r") + currNode.data;

        // 计算当前位于树的第几层
        int currLevel = ((rowIndex + 1) / 2);
        // 若到了最后一层，则返回
        if (currLevel == treeDepth) return;
        // 计算当前行到下一行，每个元素之间的间隔（下一行的列索引与当前元素的列索引之间的间隔）
        int gap = treeDepth - currLevel - 1;

        // 对左儿子进行判断，若有左儿子，则记录相应的"/"与左儿子的值
        if (currNode.leftChild != null) {
            res[rowIndex + 1][columnIndex - gap] = "/";
            writeArray(currNode.leftChild, rowIndex + 2, columnIndex - gap * 2, res, treeDepth);
        }

        // 对右儿子进行判断，若有右儿子，则记录相应的"\"与右儿子的值
        if (currNode.rightChild != null) {
            res[rowIndex + 1][columnIndex + gap] = "\\";
            writeArray(currNode.rightChild, rowIndex + 2, columnIndex + gap * 2, res, treeDepth);
        }
    }

    public void show(RBNode root) {
        if (root == null) System.out.println("EMPTY!");
        // 得到树的深度
        int treeDepth = getTreeDepth(root);

        // 最后一行的宽度为2的（n - 1）次方乘3，再加1
        // 作为整个二维数组的宽度
        int arrayHeight = treeDepth * 2 - 1;
        int arrayWidth = (2 << (treeDepth - 2)) * 3 + 1;
        // 用一个字符串数组来存储每个位置应显示的元素
        String[][] res = new String[arrayHeight][arrayWidth];
        // 对数组进行初始化，默认为一个空格
        for (int i = 0; i < arrayHeight; i ++) {
            for (int j = 0; j < arrayWidth; j ++) {
                res[i][j] = " ";
            }
        }

        // 从根节点开始，递归处理整个树
        // res[0][(arrayWidth + 1)/ 2] = (char)(root.val + '0');
        writeArray(root, 0, arrayWidth/ 2, res, treeDepth);

        // 此时，已经将所有需要显示的元素储存到了二维数组中，将其拼接并打印即可
        for (String[] line: res) {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < line.length; i ++) {
                sb.append(line[i]);
                if (line[i].length() > 1 && i <= line.length - 1) {
                    i += line[i].length() > 4 ? 2: line[i].length() - 1;
                }
            }
            System.out.println(sb.toString());
        }
    }

}
